import cv2 as cv
import numpy as np
from python_ai.common.xcommon import *
import os
import sys
import datetime

FILE_NAME = os.path.basename(__file__)
SAVE_DIR = 'd.' + FILE_NAME + '.d'
os.makedirs(SAVE_DIR, exist_ok=True)
SAVE_PATH = os.path.join(SAVE_DIR, ensure_filename(str(datetime.datetime.now())) + '.avi')

path = '../../../../large_data/video/cam/save_blue_video.avi'
cap = cv.VideoCapture(path)
W = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))  # ATTENTION cap.get
H = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
print(W, H)

H1, H2, W1, W2 = 228, 447, 176, 391
ROI_W = W2 - W1
ROI_H = H2 - H1
writerFourcc = cv.VideoWriter_fourcc(*'XVID')  # ATTENTION XVID x-伟德
wt = cv.VideoWriter(SAVE_PATH, writerFourcc, 25, (ROI_W, ROI_H), True)

blue = np.uint8((86, 116, 140))
blue = np.expand_dims(blue, (0, 1))
print(blue)
hsv_tgt = cv.cvtColor(blue, cv.COLOR_RGB2HSV)
hsv_tgt = hsv_tgt[0, 0]
lower_border = np.uint8((hsv_tgt[0] - 30, 100, 100))
upper_border = np.uint8((hsv_tgt[0] + 30, 255, 255))

kernel = np.ones((3, 3), np.uint8)

SCALE = 2
FPS = 25
INTERVAL = SCALE * 1000 // FPS
QUIT_SET = set([ord('q'), ord('Q'), 27])
first = True
# (176, 228) (391, 447)
# (176:391, 228:447)
while True:
    ret, frame = cap.read()

    if not ret:
        print('No frame, existing ...')
        break

    if first:
        first = False
        cv.imwrite(os.path.join(SAVE_DIR, 'out.jpg'), frame)

    # roi
    roi = frame[228:447, 176:391]

    cv.imshow('Video', roi)

    # filter color
    hsv = cv.cvtColor(roi, cv.COLOR_BGR2HSV)

    mask = cv.inRange(hsv, lower_border, upper_border)

    cv.imshow('mask', mask)

    # process mask
    mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel, iterations=2)
    mask = cv.morphologyEx(mask, cv.MORPH_CLOSE, kernel, iterations=2)

    cv.imshow('processed mask', mask)

    # filter image
    res = cv.bitwise_and(roi, roi, mask=mask)
    cv.imshow('filtered', res)

    wt.write(roi)

    k = cv.waitKey(INTERVAL) & 0xFF
    if k in QUIT_SET:
        break

cap.release()
wt.release()
cv.destroyAllWindows()
